#include "log_helper.hpp"
#include <helib/helib.h>

//! Including NumbTh.cpp is a very bad design, so here is a workingaround
// // #include "NumbTh.cpp"
//* Copy from NumbTh.cpp in helib (Below)
// The function compOrder(orders, classes,flag,m) computes the order of elements
// of the quotient group, relative to current equivalent classes. If flag==1
// then also check if the order is the same as in (Z/mZ)^* and store the order
// with negative sign if not.
static void compOrder(std::vector<long> &orders, std::vector<long> &classes,
                      long m) {
  NTL::mulmod_t minv = NTL::PrepMulMod(m);

  orders[0] = 0;
  orders[1] = 1;
  for (long i = 2; i < m; i++) {
    if (classes[i] <= 1) { // ignore i not in Z_m^* and order-0 elements
      orders[i] = (classes[i] == 1) ? 1 : 0;
      continue;
    }

    if (classes[i] < i) { // not a pivot
      orders[i] = orders[classes[i]];
      continue;
    }

    NTL::mulmod_precon_t iminv = NTL::PrepMulModPrecon(i, m, minv);

    // For an element i>1, the order is at least 2
    long j = NTL::MulModPrecon(i, i, m, iminv);
    long ord = 2;
    while (classes[j] != 1) {
      j = NTL::MulModPrecon(j, i, m, iminv); // next element in <i>
      ord++; // count how many steps until we reach 1
    }

    orders[i] = ord;
  }
}
// The function conjClasses(classes,g,m) unifies equivalence classes that
// have elements which are a factor of g apart, the pivot of the unified
// class is the smallest element in that class.
static void conjClasses(std::vector<long> &classes, long g, long m) {
  NTL::mulmod_t minv = NTL::PrepMulMod(m);

  for (long i = 0; i < m; i++) {
    if (classes[i] == 0)
      continue; // i \notin (Z/mZ)^*

    if (classes[i] < i) { // i is not a pivot, updated its pivot
      classes[i] = classes[classes[i]];
      continue;
    }

    // If i is a pivot, update other pivots to point to it
    NTL::mulmod_precon_t gminv = NTL::PrepMulModPrecon(g, m, minv);
    long j = NTL::MulModPrecon(i, g, m, gminv);
    while (classes[j] != i) {
      classes[classes[j]] = i; // Merge the equivalence classes of j and i

      // Note: if classes[j]!=j then classes[j] will be updated later,
      //       when we get to i=j and use the code for "i not pivot".

      j = NTL::MulModPrecon(j, g, m, gminv);
    }
  }
}
//* Copy from NumbTh.cpp in helib (Above)

void find_t(long &t, unsigned long t_bits, const long m,
            const std::vector<long> &candidates) {
  // t_bound=1 means that log_2(t) is exactly equal to t_bits+1
  const int t_bound = 1;
  static unsigned long t_bits_store = 0;
  static unsigned long t_store = 0;
  if (t_bits != t_bits_store) {
    t_bits_store = t_bits;
    t_store = 0;
  }

  std::vector<long> classes(m);
  std::vector<long> orders(m);

  // class[i] !=0 indicates that i \in (ZZ_m)^*
  for (long i = 0; i < m; i++) {
    if (NTL::GCD(i, m) != 1)
      classes[i] = 0;
    else
      classes[i] = i;
  }

  // 计算(ZZ_m)^*中每个元素的阶
  compOrder(orders, classes, m);
  long ord_p = helib::phi_N(m);
  for (auto f : candidates) {
    ord_p /= f;
  }

  bool found = false;
  static unsigned long i = 1;
  for (; i < m; i++) {
    if (classes[i] != 0 && (ord_p == orders[i])) {
      // 找到阶为 Phi(m)/(row*col) 的元素
      long temp = (t_store % m == i) ? t_store + m : i;

      while (log2(temp) < t_bits + t_bound) {
        // temp至多为t_bits+t_bound位

        while ((!NTL::ProbPrime(temp) || log2(temp) < t_bits) &&
               log2(temp) < t_bits + t_bound) {
          // 1. 增大temp
          // 2. 并保证temp在(ZZ_m)^*中的阶不变
          temp += m;
        }

        if (NTL::ProbPrime(temp) && log2(temp) >= t_bits &&
            log2(temp) < t_bits + t_bound) {
          // 1. temp为素数(此时log2(temp)不可能为整数)
          // 2. temp至少为t_bits+1位
          // 3. temp至多为t_bits+t_bound位

          //   cout << i << " " << temp << " " << orders[i] << endl;
          t = temp;
          found = true;
          t_store = t;
          break;
        }
      }

      if (found) {
        break;
      }
    }
  }
  ASSERT_PRINTF(found, "No t of order %ld found in (ZZ_m)^*, where m = %ld\n",
                ord_p, m);
}

void find_ti(long &t, unsigned long t_bits, int t_index, const long m,
             const std::vector<long> &candidates) {
  ASSERT_PRINTF(t_index >= 0, "Invalid index of t\n");

  // t_bound=1 means that log_2(t) is exactly equal to t_bits+1
  const int t_bound = 1;
  int cnt = 0;

  std::vector<long> classes(m);
  std::vector<long> orders(m);

  // class[i] !=0 indicates that i \in (ZZ_m)^*
  for (long i = 0; i < m; i++) {
    if (NTL::GCD(i, m) != 1)
      classes[i] = 0;
    else
      classes[i] = i;
  }

  // 计算(ZZ_m)^*中每个元素的阶
  compOrder(orders, classes, m);
  long ord_p = helib::phi_N(m);
  for (auto f : candidates) {
    ord_p /= f;
  }

  bool found = false;
  for (int i = 1; i < m; i++) {
    if (classes[i] != 0 && (ord_p == orders[i])) {
      // 找到阶为 Phi(m)/(row*col) 的元素
      long temp = i;

      while (log2(temp) < t_bits + t_bound) {
        // temp至多为t_bits+t_bound位

        while ((!NTL::ProbPrime(temp) || log2(temp) < t_bits) &&
               log2(temp) < t_bits + t_bound) {
          // 1. 增大temp
          // 2. 并保证temp在(ZZ_m)^*中的阶不变
          temp += m;
        }

        if (NTL::ProbPrime(temp) && log2(temp) >= t_bits &&
            log2(temp) < t_bits + t_bound) {
          if (cnt < t_index) {
            cnt++;
            temp += m;
            continue;
          }
          // 1. temp为素数(此时log2(temp)不可能为整数)
          // 2. temp至少为t_bits+1位
          // 3. temp至多为t_bits+t_bound位

          //   cout << i << " " << temp << " " << orders[i] << endl;
          t = temp;
          found = true;
          break;
        }
      }

      if (found) {
        break;
      }
    }
  }
  ASSERT_PRINTF(found, "No t of order %ld found in (ZZ_m)^*, where m =%ld\n ",
                ord_p, m);
}

void find_first_t(long &t, unsigned long t_bits, const long m,
                  const std::vector<long> &candidates) {
  find_ti(t, t_bits, 0, m, candidates);
}

void find_ti_auto(long &t, unsigned long t_bits, const long m,
                  const std::vector<long> &candidates) {
  std::mutex mtx;

  // t_bound=1 means that log_2(t) is exactly equal to t_bits+1
  const int t_bound = 1;
  mtx.lock();
  static int static_t_index = -1;
  static_t_index++;
  int t_index = static_t_index;
  mtx.unlock();
  int cnt = 0;

  std::vector<long> classes(m);
  std::vector<long> orders(m);

  // class[i] !=0 indicates that i \in (ZZ_m)^*
  for (long i = 0; i < m; i++) {
    if (NTL::GCD(i, m) != 1)
      classes[i] = 0;
    else
      classes[i] = i;
  }

  // 计算(ZZ_m)^*中每个元素的阶
  compOrder(orders, classes, m);
  long ord_p = helib::phi_N(m);
  for (auto f : candidates) {
    ord_p /= f;
  }

  bool found = false;
  for (int i = 1; i < m; i++) {
    if (classes[i] != 0 && (ord_p == orders[i])) {
      // 找到阶为 Phi(m)/(row*col) 的元素
      long temp = i;

      while (log2(temp) < t_bits + t_bound) {
        // temp至多为t_bits+t_bound位

        while ((!NTL::ProbPrime(temp) || log2(temp) < t_bits) &&
               log2(temp) < t_bits + t_bound) {
          // 1. 增大temp
          // 2. 并保证temp在(ZZ_m)^*中的阶不变
          temp += m;
        }

        if (NTL::ProbPrime(temp) && log2(temp) >= t_bits &&
            log2(temp) < t_bits + t_bound) {
          if (cnt < t_index) {
            cnt++;
            temp += m;
            continue;
          }
          // 1. temp为素数(此时log2(temp)不可能为整数)
          // 2. temp至少为t_bits+1位
          // 3. temp至多为t_bits+t_bound位

          //   cout << i << " " << temp << " " << orders[i] << endl;
          t = temp;
          found = true;
          break;
        }
      }

      if (found) {
        break;
      }
    }
  }
  ASSERT_PRINTF(found, "No t of order %ld found in (ZZ_m)^*, where m = %ld\n ",
                ord_p, m);
}

void find_generators(std::vector<long> &gens, std::vector<long> &ords, long m,
                     long p, const std::vector<long> &candidates) {
  gens.clear();
  ords.clear();
  std::vector<long> classes(m);
  std::vector<long> orders(m);
  for (long i = 0; i < m; i++) {
    if (NTL::GCD(i, m) != 1)
      classes[i] = 0;
    else
      classes[i] = i;
  }
  // 计算 (ZZ_m)^* /<p> 的等价类(代表元选择最小元素)
  conjClasses(classes, p % m, m);

  // candidates = {row, col}:
  // ord(g1) in (ZZ_m)^* /<p> = row,
  // ord(g2) in (ZZ_m)^* /<p, g1> = col
  // 1. 计算 (ZZ_m)^* /<p> 中元素的阶
  // 2. 计算 (ZZ_m)^* /<p, g1> 的等价类(代表元选择最小元素)
  // 3. 计算 (ZZ_m)^* /<p, g1> 中元素的阶
  for (int i = 0; i < helib::lsize(candidates); i++) {
    compOrder(orders, classes, m);
    for (int j = 0; j < m; j++)
      if (orders[j] == candidates[i]) {
        // FIXME: Can we always find "good dimensions" ?
        long val = NTL::PowerMod(j, orders[j], m);
        if (val != 1l) {
          //! Skip bad dimension
          continue;
        }
        gens.push_back(j);
        ords.push_back(candidates[i]);
        conjClasses(classes, j, m);
        break;
      }
  }
  ASSERT_PRINTF(gens.size() == candidates.size(),
                "Can not find generators s.t. ord(g) = candidates\n");
  return;
}